import torch.cuda
from torch import nn

from models.MLP import MLP

device = 'cuda' if torch.cuda.is_available() else 'cpu'
import torch.nn.functional as F

class LabelModel(nn.Module):
    """
     每个属性一个分类器，
     损失函数使用BCELoss
    """
    def __init__(self, dataset, args):
        super(LabelModel, self).__init__()
        # 输入维度为图片特征维度，输出维度为属性的数量
        self.img_classifier = MLP(len(dataset.attrs[0]), len(dataset.classes), relu=False).to(device)

    def forward(self, imgs, labels=None, attrs=None):
        label_preds = self.img_classifier(imgs)
        return label_preds
